from typing import Type, Callable, Optional
from dataclasses import dataclass
from injector import inject
from langchain_core.tools import BaseTool
from pydantic.v1 import BaseModel, create_model, Field
from internal.core.tools.api_tools.entities.api_tools_bean import ParameterIn, ParameterTypeMap
from requests import request
from langchain_core.tools import StructuredTool

from internal.core.tools.api_tools.entities import ToolEntity


@inject
@dataclass
class ApiProviderManager:
    @classmethod
    def _create_tool_func_from_tool_entity(cls, tool_entity: ToolEntity) -> Callable:
        def tool_fun(**kwargs):
            parmpters = {
                ParameterIn.PATH: {},
                ParameterIn.QUERY: {},
                ParameterIn.COOKIE: {},
                ParameterIn.REQUEST_BODY: {}
            }
            parmpters_map = {par.get("name"): par for par in tool_entity.parameters}
            parmpters_headers = {k: v for k, v in enumerate(tool_entity.headers)}
            for k, v in kwargs.items():
                parmpt = parmpters_map.get(k)
                if parmpt is None:
                    continue
                parmpters[parmpt.get("in", ParameterIn.QUERY)][k] = v
            print(parmpters[ParameterIn.QUERY])
            return request(
                method=tool_entity.method,
                url=tool_entity.url.format(**parmpters[ParameterIn.PATH]),
                params=parmpters[ParameterIn.QUERY],
                json=parmpters[ParameterIn.REQUEST_BODY],
                # headers={**parmpters_headers, **parmpters[ParameterIn.HEADER]},
                cookies=parmpters[ParameterIn.COOKIE]
            ).text

        return tool_fun

    @classmethod
    def _create_model_from_parameters(cls, parameters: list[dict]) -> Type[BaseModel]:
        files = {}
        for parameter in parameters:
            file_name = parameter.get("name")
            file_type = ParameterTypeMap.get(parameter.get("type"), str)
            file_required = parameter.get("required", True)
            fild_description = parameter.get("description")
            files[file_name] = (
                file_type if file_required else Optional[file_type],
                Field(description=fild_description)
            )
        return create_model("CustomModel", **files)

    def get_tool(self, tool_entity: ToolEntity) -> BaseTool:
        tools = StructuredTool.from_function(
            func=self._create_tool_func_from_tool_entity(tool_entity),
            name=tool_entity.name,
            description=tool_entity.description,
            args_schema=self._create_model_from_parameters(tool_entity.parameters)
        )
        return tools
